#include "stun_protocol_attr.h"

#include "./../compress/compress.h"
#include "./../string_util/string_util.h"

#include <arpa/inet.h>
#include <cstdio>
#include <cstring>
#include <netinet/in.h>
#include <sys/socket.h>

std::shared_ptr<StunAttr> StunAttr::CreateStunAttr(int attr_type)
{
    std::shared_ptr<StunAttr> attr = nullptr;
    switch (attr_type)
    {
        case EmMappedAddress:
        {
            attr = std::make_shared<StunMappedAddressAttr>();
            break;
        }
        case EmXorMappedAddress:
        {
            attr = std::make_shared<StunXorMappedAddressAttr>();
            break;
        }
        case EmRequestTransport:
        {
            attr = std::make_shared<StunRequestTransportAttr>();
            break;
        }
        case EmUserName:
        {
            attr = std::make_shared<StunUserNameAttr>();
            break;
        }
        case EmFingerPrint:
        {
            attr = std::make_shared<StunFingerPrintAttr>();
            break;
        }
        default:
            break;
    }
    return attr;
}

StunAttr::StunAttr()
{
    m_attr_type    = 0;
    m_attr_length  = 0;
    m_attr_padding = 0;
}

StunAttr::StunAttr(int type) : m_attr_type(type)
{
    m_attr_length  = 0;
    m_attr_padding = 0;
}

StunAttr::StunAttr(int type, int length) : m_attr_type(type), m_attr_length(length)
{
    m_attr_padding = 0;
}

StunAttr::~StunAttr() = default;

int StunAttr::AttrType() const
{
    return m_attr_type;
}

int StunAttr::AttrLength() const
{
    return m_attr_length;
}

int StunAttr::AttrPadding() const
{
    return m_attr_padding;
}

uint32_t StunAttr::MagicCookie() const
{
    return m_magic_cookie;
}

void StunAttr::AttrLength(int length)
{
    m_attr_length = length;
}

void StunAttr::MagicCookie(uint32_t cookie)
{
    m_magic_cookie = cookie;
}

void StunAttr::AttrPadding(int length)
{
    m_attr_padding = length;
}

void StunAttr::TransactionID(uint8_t * data, int length)
{
    m_transaction_id.resize(length);
    memcpy(m_transaction_id.data(), data, length);
}

uint8_t * StunAttr::TransactionID()
{
    return m_transaction_id.data();
}

int StunAttr::ComposeBaseAttrToBuffer(uint8_t * data, int data_length) const
{
    uint16_t type  = StringUtil::HtoBe16(m_attr_type);
    uint16_t len   = StringUtil::HtoBe16(m_attr_length);
    int      index = 0;
    memcpy(data, &type, sizeof(type));
    index += 2;
    memcpy(data + index, &len, sizeof(len));
    index += 2;
    return index;
}

void StunAttr::DumpAttr()
{
    printf("type:0x%04x,length:%d,padding:%d.\n", m_attr_type, m_attr_length, m_attr_padding);
}

StunMappedAddressAttr::StunMappedAddressAttr() : StunAttr(EmMappedAddress)
{
}

StunMappedAddressAttr::~StunMappedAddressAttr() = default;

void StunMappedAddressAttr::DumpAttr()
{
    StunAttr::DumpAttr();
    printf("family:0x%02x,port:%d,ip：%s.\n", m_family_type, HostPort(), HostIP().c_str());
}

int StunMappedAddressAttr::ComposeAttrToBuffer(uint8_t * data, int length)
{
    int index       = ComposeBaseAttrToBuffer(data, 0);
    *(data + index) = 0x00;
    index          += 1;
    *(data + index) = m_family_type;
    index          += 1;
    if (m_family_type == 0x01)
    {
        memcpy(data + index, &m_address.sin_port, sizeof(m_address.sin_port));
        index += sizeof(m_address.sin_port);
        memcpy(data + index, &m_address.sin_addr.s_addr, sizeof(m_address.sin_addr.s_addr));
        index += sizeof(m_address.sin_addr.s_addr);
    }
    else
    {
        memcpy(data + index, &m_address_v6.sin6_port, sizeof(m_address_v6.sin6_port));
        index += sizeof(m_address_v6.sin6_port);
        memcpy(data + index, &m_address_v6.sin6_addr, sizeof(m_address_v6.sin6_addr));
        index += sizeof(m_address_v6.sin6_addr);
    }
    return index;
}

int StunMappedAddressAttr::ParseBufferToAttr(const uint8_t * data, int data_length, int & use_length)
{
    int index     = 1; // 1byte resvered
    m_family_type = *(data + index);
    index        += 1;
    uint16_t port = 0;
    memcpy(&port, data + index, sizeof(port));
    index += sizeof(port);
    switch (m_family_type)
    {
        case 0x01:
        {
            memcpy(&m_address.sin_addr.s_addr, data + index, sizeof(m_address.sin_addr.s_addr));
            m_address.sin_port   = port;
            m_address.sin_family = AF_INET;
            index               += sizeof(m_address.sin_addr);
            break;
        }
        case 0x02:
        {
            memcpy(&m_address_v6.sin6_addr, data + index, sizeof(m_address_v6.sin6_addr));
            m_address_v6.sin6_family = AF_INET6;
            m_address_v6.sin6_port   = port;
            index                   += sizeof(m_address_v6.sin6_addr);
            break;
        }
        default:
            break;
    }
    use_length += index;
    return EmParseSuccess;
}

uint16_t StunMappedAddressAttr::HostPort() const
{
    if (m_family_type == 0x01)
    {
        return ntohs(m_address.sin_port);
    }
    else
    {
        return ntohs(m_address_v6.sin6_port);
    }
}

std::string StunMappedAddressAttr::HostIP()
{
    char ip[129] = {};
    if (m_family_type == 0x01)
    {
        inet_ntop(m_address.sin_family, &m_address.sin_addr.s_addr, ip, sizeof(ip) - 1);
    }
    else
    {
        inet_ntop(m_address_v6.sin6_family, &m_address_v6.sin6_addr, ip, sizeof(ip) - 1);
    }
    return ip;
}

void StunMappedAddressAttr::AddressInfo(uint8_t family, uint16_t port, const std::string & ip)
{
    m_family_type = family;
    if (family == 0x01)
    {
        m_address.sin_family = AF_INET;
        m_address.sin_port   = htons(port);
        inet_pton(AF_INET, ip.c_str(), &m_address.sin_addr.s_addr);
        AttrLength(8);
    }
    else
    {
        m_address_v6.sin6_family = AF_INET6;
        m_address.sin_port       = htons(port);
        inet_pton(AF_INET6, ip.c_str(), &m_address_v6.sin6_addr);
        AttrLength(20);
    }
}

StunXorMappedAddressAttr::StunXorMappedAddressAttr() : StunAttr(EmXorMappedAddress)
{
}

StunXorMappedAddressAttr::~StunXorMappedAddressAttr() = default;

uint16_t StunXorMappedAddressAttr::HostPort() const
{
    if (m_family == 0x01)
    {
        return ntohs(m_address.sin_port);
    }
    else
    {
        return ntohs(m_address_v6.sin6_port);
    }
}

std::string StunXorMappedAddressAttr::HostIP()
{
    char ip[129] = {};
    if (m_family == 0x01)
    {
        inet_ntop(m_address.sin_family, &m_address.sin_addr, ip, sizeof(ip) - 1);
    }
    else
    {
        inet_ntop(m_address_v6.sin6_family, &m_address_v6.sin6_addr, ip, sizeof(ip) - 1);
    }
    return ip;
}

int StunXorMappedAddressAttr::ComposeAttrToBuffer(uint8_t * data, int length)
{
    int index = ComposeBaseAttrToBuffer(data, length);
    if (m_family == 0x01)
    {
        uint16_t port   = StringUtil::HtoBe16(ntohs(m_address.sin_port) ^ (MagicCookie() >> 16));
        uint32_t addr   = m_address.sin_addr.s_addr ^ StringUtil::HtoBe32(MagicCookie());
        *(data + index) = 0x00;
        index          += 1;
        *(data + index) = 0x01;
        index          += 1;
        memcpy(data + index, &port, sizeof(port));
        index += sizeof(port);
        memcpy(data + index, &addr, sizeof(addr));
        index += sizeof(addr);
    }
    else
    {
        uint16_t  port   = StringUtil::HtoBe16(ntohs(m_address_v6.sin6_port) ^ (MagicCookie() >> 16));
        auto      addr   = m_address_v6.sin6_addr;
        uint32_t  cookie = MagicCookie();
        uint8_t * ID     = TransactionID();
        uint8_t * ptr    = reinterpret_cast<uint8_t *>(&addr);
        *(ptr)           = *(ptr) ^ ((cookie >> 24) & 0xff);
        *(ptr + 1)       = *(ptr + 1) ^ ((cookie >> 16) & 0xff);
        *(ptr + 2)       = *(ptr + 2) ^ ((cookie >> 8) & 0xff);
        *(ptr + 3)       = *(ptr + 3) ^ (cookie & 0xff);
        for (int i = 0; i < 12; i++)
        {
            *(ptr + i + 4) = *(ptr + i + 4) ^ *(ID + i);
        }
        *(data + index) = 0x00;
        index          += 1;
        *(data + index) = 0x02;
        index          += 1;
        memcpy(data + index, &port, sizeof(port));
        index += 2;
        memcpy(data + index, &addr, sizeof(addr));
        index += sizeof(addr);
    }
    return index;
}

void StunXorMappedAddressAttr::DumpAttr()
{
    StunAttr::DumpAttr();
    printf("family:0x%02x,port:%d,ip:%s.\n", m_family, HostPort(), HostIP().c_str());
}

int StunXorMappedAddressAttr::ParseBufferToAttr(const uint8_t * data, int data_length, int & use_length)
{
    if (data_length < AttrLength())
    {
        return EmParseNeedData;
    }
    int     index  = 1; // resverd 1 byte
    uint8_t family = *(data + index);
    if (family != 0x01 && family != 0x02)
    {
        return EmParseFailed;
    }
    m_family      = family;
    index         = index + 1;
    uint16_t port = 0;
    memcpy(&port, data + index, sizeof(port));
    port                 = StringUtil::BeToH16(port);
    port                 = port ^ (MagicCookie() >> 16);
    index                = index + sizeof(port);
    switch (family)
    {
        case 0x01:
        {
            memcpy(&m_address.sin_addr.s_addr, data + index, sizeof(m_address.sin_addr.s_addr));
            m_address.sin_addr.s_addr = m_address.sin_addr.s_addr ^ (StringUtil::BeToH32(MagicCookie()));
            m_address.sin_family      = AF_INET;
            m_address.sin_port        = htons(port);
            index                    += sizeof(m_address.sin_addr.s_addr);
            break;
        }
        case 0x02:
        {
            memcpy(&m_address_v6.sin6_addr, data + index, sizeof(m_address_v6.sin6_addr));
            uint32_t  cookie = MagicCookie();
            uint8_t * ID     = TransactionID();
            uint8_t * ptr    = reinterpret_cast<uint8_t *>(&m_address_v6.sin6_addr);
            *(ptr)           = *(ptr) ^ ((cookie >> 24) & 0xff);
            *(ptr + 1)       = *(ptr + 1) ^ ((cookie >> 16) & 0xff);
            *(ptr + 2)       = *(ptr + 2) ^ ((cookie >> 8) & 0xff);
            *(ptr + 3)       = *(ptr + 3) ^ (cookie & 0xff);
            for (int i = 0; i < 12; i++)
            {
                *(ptr + i + 4) = *(ptr + i + 4) ^ *(ID + i);
            }
            m_address_v6.sin6_family = AF_INET6;
            m_address_v6.sin6_port   = htons(port);
            index                   += sizeof(m_address_v6.sin6_addr);
            break;
        }
    }
    use_length += index;
    return EmParseSuccess;
}

void StunXorMappedAddressAttr::AddressInfo(uint8_t family, uint16_t port, const std::string & ip)
{
    m_family = family;
    if (family == 0x01)
    {
        m_address.sin_family = AF_INET;
        m_address.sin_port   = htons(port);
        inet_pton(AF_INET, ip.c_str(), &m_address.sin_addr.s_addr);
        AttrLength(8);
    }
    else
    {
        m_address_v6.sin6_family = AF_INET6;
        m_address_v6.sin6_port   = htons(port);
        inet_pton(AF_INET6, ip.c_str(), &m_address_v6.sin6_addr);
        AttrLength(20);
    }
}

StunRequestTransportAttr::StunRequestTransportAttr() : StunAttr(EmRequestTransport)
{
}

StunRequestTransportAttr::~StunRequestTransportAttr()
= default;

int StunRequestTransportAttr::ParseBufferToAttr(const uint8_t * data, int data_length, int & use_length)
{
    int index        = 0;
    m_transport_type = *(data + index);
    index           += 1;
    index           += 3;
    use_length      += index;
    return EmParseSuccess;
}

int StunRequestTransportAttr::ComposeAttrToBuffer(uint8_t * data, int data_length)
{
    int index           = ComposeBaseAttrToBuffer(data, data_length);
    *(data + index)     = m_transport_type;
    *(data + index + 1) = 0;
    *(data + index + 2) = 0;
    *(data + index + 3) = 0;
    index              += 4;
    return index;
}

void StunRequestTransportAttr::DumpAttr()
{
    StunAttr::DumpAttr();
    printf("transport:%s.\n", m_transport_type == 0x11 ? "UDP" : "TCP");
}

void StunRequestTransportAttr::Transport(uint8_t type)
{
    m_transport_type = type;
    AttrLength(4);
}

int StunRequestTransportAttr::Transport() const
{
    return m_transport_type;
}

StunUserNameAttr::StunUserNameAttr() : StunAttr(EmUserName)
{
}

StunUserNameAttr::~StunUserNameAttr()
{
    m_username.clear();
}

void StunUserNameAttr::DumpAttr()
{
    StunAttr::DumpAttr();
    printf("username:%s.\n", m_username.c_str());
}

int StunUserNameAttr::ComposeAttrToBuffer(uint8_t * data, int data_length)
{
    int index   = ComposeBaseAttrToBuffer(data, data_length);
    int padding = AttrPadding();
    memcpy(data + index, m_username.c_str(), m_username.size());
    index += static_cast<int>(m_username.size());
    for (int i = 0; i < padding; i++)
    {
        *(data + index + i) = 0x00;
    }
    index += padding;
    return index;
}

int StunUserNameAttr::ParseBufferToAttr(const uint8_t * data, int data_length, int & use_length)
{
    int attr_length = AttrLength();
    int padding     = StunAttr::padding_len - (attr_length % StunAttr::padding_len);
    AttrPadding(padding);
    m_username.assign((const char *)(data), attr_length);
    use_length += attr_length;
    return EmParseSuccess;
}

void StunUserNameAttr::UserName(const std::string & user_name)
{
    m_username  = user_name;
    int padding = static_cast<int>(StunAttr::padding_len - user_name.size() % StunAttr::padding_len);
    AttrPadding(padding);
    AttrLength(static_cast<int>(m_username.size()));
}

std::string StunUserNameAttr::UserName()
{
    return m_username;
}

StunFingerPrintAttr::StunFingerPrintAttr() : StunAttr(EmFingerPrint)
{
    m_crc32       = 0;
    m_fingerprint = 0;
}

StunFingerPrintAttr::~StunFingerPrintAttr()
{
    m_crc32 = 0;
}

int StunFingerPrintAttr::ParseBufferToAttr(const uint8_t * data, int data_length, int & use_length)
{
    memcpy(&m_fingerprint, data, sizeof(m_fingerprint));
    m_fingerprint = StringUtil::BeToH32(m_fingerprint);
    m_crc32       = m_fingerprint ^ 0x5354554e;
    return EmParseSuccess;
}

int StunFingerPrintAttr::ComposeAttrToBuffer(uint8_t * data, int data_length)
{
    int      index       = ComposeBaseAttrToBuffer(data, data_length);
    uint32_t fingerprint = StringUtil::HtoBe32(m_fingerprint);
    memcpy(data + index, &fingerprint, sizeof(fingerprint));
    index += sizeof(m_fingerprint);
    return index;
}

void StunFingerPrintAttr::DumpAttr()
{
    StunAttr::DumpAttr();
    printf("fingerprint:0x%08x, crc32:0x%08x.\n", m_fingerprint, m_crc32);
}

uint32_t StunFingerPrintAttr::FingerPrint() const
{
    return m_fingerprint;
}

uint32_t StunFingerPrintAttr::ComputeFingerPrint(const uint8_t * data, int length)
{
    CompressData c     = {};
    uint32_t     crc32 = c.CRC32(data, length);
    m_crc32            = crc32;
    m_fingerprint      = m_crc32 ^ 0x5354554e;
    return m_crc32;
}
